{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2f5db563",
   "metadata": {},
   "source": [
    "## 导入模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b986dd5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from qdrant_client import models, QdrantClient\n",
    "from sentence_transformers import SentenceTransformer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f714e2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "encoder = SentenceTransformer(\"all-MiniLM-L6-v2\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf0cc092",
   "metadata": {},
   "source": [
    "## 数据集"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1aafb9c",
   "metadata": {},
   "outputs": [],
   "source": [
    "documents = [\n",
    "    {\n",
    "        \"name\": \"The Time Machine\",\n",
    "        \"description\": \"A man travels through time and witnesses the evolution of humanity.\",\n",
    "        \"author\": \"H.G. Wells\",\n",
    "        \"year\": 1895,\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"Ender's Game\",\n",
    "        \"description\": \"A young boy is trained to become a military leader in a war against an alien race.\",\n",
    "        \"author\": \"Orson Scott Card\",\n",
    "        \"year\": 1985,\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"Brave New World\",\n",
    "        \"description\": \"A dystopian society where people are genetically engineered and conditioned to conform to a strict social hierarchy.\",\n",
    "        \"author\": \"Aldous Huxley\",\n",
    "        \"year\": 1932,\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"The Hitchhiker's Guide to the Galaxy\",\n",
    "        \"description\": \"A comedic science fiction series following the misadventures of an unwitting human and his alien friend.\",\n",
    "        \"author\": \"Douglas Adams\",\n",
    "        \"year\": 1979,\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"Dune\",\n",
    "        \"description\": \"A desert planet is the site of political intrigue and power struggles.\",\n",
    "        \"author\": \"Frank Herbert\",\n",
    "        \"year\": 1965,\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"Foundation\",\n",
    "        \"description\": \"A mathematician develops a science to predict the future of humanity and works to save civilization from collapse.\",\n",
    "        \"author\": \"Isaac Asimov\",\n",
    "        \"year\": 1951,\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"Snow Crash\",\n",
    "        \"description\": \"A futuristic world where the internet has evolved into a virtual reality metaverse.\",\n",
    "        \"author\": \"Neal Stephenson\",\n",
    "        \"year\": 1992,\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"Neuromancer\",\n",
    "        \"description\": \"A hacker is hired to pull off a near-impossible hack and gets pulled into a web of intrigue.\",\n",
    "        \"author\": \"William Gibson\",\n",
    "        \"year\": 1984,\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"The War of the Worlds\",\n",
    "        \"description\": \"A Martian invasion of Earth throws humanity into chaos.\",\n",
    "        \"author\": \"H.G. Wells\",\n",
    "        \"year\": 1898,\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"The Hunger Games\",\n",
    "        \"description\": \"A dystopian society where teenagers are forced to fight to the death in a televised spectacle.\",\n",
    "        \"author\": \"Suzanne Collins\",\n",
    "        \"year\": 2008,\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"The Andromeda Strain\",\n",
    "        \"description\": \"A deadly virus from outer space threatens to wipe out humanity.\",\n",
    "        \"author\": \"Michael Crichton\",\n",
    "        \"year\": 1969,\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"The Left Hand of Darkness\",\n",
    "        \"description\": \"A human ambassador is sent to a planet where the inhabitants are genderless and can change gender at will.\",\n",
    "        \"author\": \"Ursula K. Le Guin\",\n",
    "        \"year\": 1969,\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"The Three-Body Problem\",\n",
    "        \"description\": \"Humans encounter an alien civilization that lives in a dying system.\",\n",
    "        \"author\": \"Liu Cixin\",\n",
    "        \"year\": 2008,\n",
    "    },\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47af220a",
   "metadata": {},
   "source": [
    "## 定义临时存储位置"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f9af076",
   "metadata": {},
   "outputs": [],
   "source": [
    "client = QdrantClient(path=\"D:/YXT_Project/vector_db_practices/qdrant/data\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9d4ba2a",
   "metadata": {},
   "source": [
    "## 创建集合"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96b738fe",
   "metadata": {},
   "source": [
    "### 添加判断，集合存在则跳过创建步骤"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a7e8214",
   "metadata": {},
   "outputs": [],
   "source": [
    "existing_collections = [col.name for col in client.get_collections().collections]\n",
    "\n",
    "if \"my_books\" not in existing_collections:\n",
    "    client.create_collection(\n",
    "        collection_name=\"my_books\",\n",
    "        vectors_config=models.VectorParams(\n",
    "            size=encoder.get_sentence_embedding_dimension(),\n",
    "            distance=models.Distance.COSINE,\n",
    "        ),\n",
    "    )\n",
    "else:\n",
    "    print(\"Collection 'my_books' already exists. Skipping creation.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6f13a27",
   "metadata": {},
   "source": [
    "## 将数据上传到集合"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9920cd41",
   "metadata": {},
   "outputs": [],
   "source": [
    "client.upload_points(\n",
    "    collection_name=\"my_books\",\n",
    "    points=[\n",
    "        models.PointStruct(\n",
    "            id=idx, vector=encoder.encode(doc[\"description\"]).tolist(), payload=doc\n",
    "        )\n",
    "        for idx, doc in enumerate(documents)\n",
    "    ],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06a18819",
   "metadata": {},
   "source": [
    "## 提问"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43c9ed7b",
   "metadata": {},
   "source": [
    "- 一轮提问"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d789d2b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "hits = client.query_points(\n",
    "    collection_name=\"my_books\",\n",
    "    query=encoder.encode(\"alien invasion\").tolist(),\n",
    "    limit=3,\n",
    ").points\n",
    "\n",
    "for hit in hits:\n",
    "    print(hit.payload, \"score:\", hit.score)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b466cd87",
   "metadata": {},
   "source": [
    "- 缩小范围提问"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0429f4a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "hits = client.query_points(\n",
    "    collection_name=\"my_books\",\n",
    "    query=encoder.encode(\"alien invasion\").tolist(),\n",
    "    query_filter=models.Filter(\n",
    "        must=[models.FieldCondition(key=\"year\", range=models.Range(gte=2000))]\n",
    "    ),\n",
    "    limit=1,\n",
    ").points\n",
    "\n",
    "for hit in hits:\n",
    "    print(hit.payload, \"score:\", hit.score)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33c8983c",
   "metadata": {},
   "source": [
    "## 查询全部的数据及向量表示"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65f810f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "scroll_result = client.scroll(\n",
    "    collection_name=\"my_books\",\n",
    "    scroll_filter=None,  # 不加过滤条件\n",
    "    with_payload=True,   # 返回 payload（文本）\n",
    "    with_vectors=True,   # 返回向量\n",
    "    limit=100,           # 一次最多返回 100 条记录\n",
    ")\n",
    "\n",
    "print(\"=== 向量库内容如下 ===\")\n",
    "for point in scroll_result[0]:  # scroll_result 是一个二元组 (points, next_page)\n",
    "    print(f\"ID: {point.id}\")\n",
    "    print(f\"Description: {point.payload.get('description')}\")\n",
    "    print(f\"Vector (前5维): {point.vector[:5]} ...\")  # 只打印前5维向量简洁展示\n",
    "    print(\"-\" * 50)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87d73c06",
   "metadata": {},
   "source": [
    "## 向量维度"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "547bf1d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"向量维度:\", encoder.get_sentence_embedding_dimension())"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py310",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
